from Index.Fractal import Fractal
from Index.MACDPoint import MACD
from DataControl.getDataFromTushare.tudata import pullDataFromTushare
import pandas as pd
import numpy as np
import Index.utils as utils
###

def checkisnan(df, f):
    a = df.values
    pos = 0
    for i in range(a.shape[0]):
        if a[i] != np.nan:
            pos = i
            break
    for i in range(pos, a.shape[0], 1):
        if i != 0 and a[i] == np.nan:
            print(f)
            print(df)
            return i

    return -1


class DataSet(object):

    """
    DataSet 类用于处理和计算金融市场数据的各种技术指标。
    该类从数据库中拉取数据，并使用这些数据计算多个常见的技术指标。它将处理后的数据存储在一个 Pandas DataFrame 中，便于后续分析和建模。
    """
    def __init__(self, ts_code, start_date=None, end_date=None): # 时间格式：yyyymmdd, 字符串
        self.pull = pullDataFromTushare()
        if start_date == None and end_date == None:
            self.data = self.pull.getallindex_daily_from_database(ts_code)
        else :
            assert(start_date != None and end_date != None)
            if utils.time_delta(start_date, end_date) < 60:
                print(utils.time_delta(start_date, end_date))
               # print('至少需要60天')
                raise Exception('至少需要60天')
            self.data = self.pull.getindex_daily_from_database(ts_code=ts_code, start=start_date, end=end_date)

        self.macd = MACD(self.data)
        self.data = self.macd.df

        self.fractal = Fractal(ts_code)
        self.fractal.raw_data.drop(index = self.fractal.raw_data.index[:self.fractal.raw_data.shape[0] - self.data.shape[0]], inplace=True)
        self.data['open-close'] = (self.data['close'] - self.data['open']) / self.data['open']
        self.data['high-low'] = (self.data['high'] - self.data['low']) / self.data['low']
        self.data['Fmark'] = self.fractal.raw_data['Fmark'].values # 分型的特征，1为顶分型， -1为底分型
       # print(self.fractal.raw_data['Fmark'])

        # print(self.data.columns)
        self.data = self.data.values
        #print(self.data)

        # 设置分型中间的值，上升2，下降-2

        """
        Fmark在这个时候要设置为最后一列
        """
        flag = 0
        for i in range(self.data.shape[0]):
            if self.data[i][-1] != 0:
                flag = self.data[i][-1]
                #flag = self.data[i][-1] + abs(self.data[i][-1]) // self.data[i][-1]
                continue
            if flag == -1:
                self.data[i][-1] = 2
            if flag == 1:
                self.data[i][-1] = -2
        flag = 0
        for i in range(self.data.shape[0] - 1, -1, -1):
            if self.data[i][-1] == -1 or self.data[i][-1] == 1:
                flag = self.data[i][-1]
                continue
            if flag == -1:
                self.data[i][-1] = -2
            if flag == 1:
                self.data[i][-1] = 2
           # print(i, self.data[i][-1])
        i = 1
        while i < self.data.shape[0] - 1:
            if self.data[i][-1] == 1 or self.data[i][-1] == -1:
                self.data[i + 1, -1] = self.data[i - 1, -1] = self.data[i, -1]
                i += 1
            i += 1


        self.data = pd.DataFrame(self.data, columns=['trade_date', 'close', 'open', 'high', 'low', 'pre_close', 'change',
                                                    'pct_chg', 'vol', 'amount', 'turnover_rate', 'RSI', 'MACD', 'DEA',
                                                    'Histogram', 'MA5', 'MA10', 'MOM', 'EMA12', 'EMA26',
                                                    'open-close', 'high-low', 'Fmark'])

        #print(self.data)

        self.data['+DI'], self.data['-DI'], self.data['ADX'], self.data['ADXR'] = self.DMI()
        # checkisnan(self.data['+DI'], '+DI')
        # checkisnan(self.data['-DI'], '-DI')
        # checkisnan(self.data['ADX'], 'ADX')
        # checkisnan(self.data['ADXR'], 'ADXR')

        self.data['TRIX'], self.data['TRMA'] = self.TRIX()
        #checkisnan(self.data['TRIX'], 'TRIX')
       # checkisnan(self.data['TRMA'], 'TRMA')

        self.data['DMA'] = self.DMA()
      #  checkisnan(self.data['DMA'], 'DMA')

        self.data['MBOLL'], self.data['UBOLL'], self.data['DBOLL'] = self.BOLL()
        #checkisnan(self.data['MBOLL'], 'MBOLL')
       # checkisnan(self.data['UBOLL'], 'UBOLL')
       # checkisnan(self.data['DBOLL'], 'DBOLL')

        self.data['W_R'] = self.W_R()
      #  checkisnan(self.data['W_R'], 'W_R')

        self.data['BIAS'] = self.BIAS()
    #    checkisnan(self.data['BIAS'], 'BIAS')

        self.data['CR'] = self.CR()
      #  checkisnan(self.data['CR'], 'CR')

        self.data['VR'] = self.VR()
      #  checkisnan(self.data['VR'], 'VR')

        self.data['OBV'] = self.OBV()
        checkisnan(self.data['OBV'], 'OBV')

        self.data[['TYP', 'WR', 'MR', 'SR', 'WS', 'MS', 'SS']] = self.MIKE()
      #  checkisnan(self.data['TYP'], 'TYP')
      #  checkisnan(self.data['WR'], 'WR')
      #  checkisnan(self.data['MR'], 'MR')
       # checkisnan(self.data['SR'], 'SR')
       # checkisnan(self.data['WS'], 'WS')
       ## checkisnan(self.data['MS'], 'MS')
        #checkisnan(self.data['SS'], 'SS')

        self.data['PVI'] = self.PVI()
       # checkisnan(self.data['PVI'], 'PVI')

        self.data['ASI'] = self.ASI()
        checkisnan(self.data['ASI'], 'ASI')

        self.data[['KDJ_K', 'KDJ_D', 'KDJ_J']] = self.KDJ()
        # checkisnan(self.data['KDJ_K'], 'KDJ_K')
        # checkisnan(self.data['KDJ_D'], 'KDJ_D')
        # checkisnan(self.data['KDJ_J'], 'KDJ_J')

        self.data['ROC'] = self.ROC()
       # checkisnan(self.data['ROC'], 'ROC')

        self.data['MTM'] = self.MTM()
       # checkisnan(self.data['MTM'], 'MTM')

        self.data[['AR', 'BR']] = self.ARBR()
      #  checkisnan(self.data['AR'], 'AR')
      #  checkisnan(self.data['BR'], 'BR')

        self.data['PSY'] = self.PSY()
       # checkisnan(self.data['PSY'], 'PSY')

        self.data['SAR'] = self.SAR()
       # checkisnan(self.data['SAR'], 'SAR')

        self.data['OBOS'] = self.OBOS()
       # checkisnan(self.data['OBOS'], 'OBOS')

        self.data['CCI'] = self.CCI()
       # checkisnan(self.data['CCI'], 'CCI')

        self.data['WVAD'] = self.WVAD()
        #checkisnan(self.data['WVAD'], 'WVAD')

        self.data['TAPI'] = self.TAPI()
        #checkisnan(self.data['TAPI'], 'TAPI')

        self.data.dropna(inplace=True)

    #def stacked_data(self):


    def DMI(self, period=14):
        #res = np.full(self.data.shape[0], np.nan)
        high = self.data['high']
        low = self.data['low']
        close = self.data['close']

        # 计算TR
        tr1 = high - low
        tr2 = abs(high - close.shift())
        tr3 = abs(low - close.shift())
        tr = pd.concat([tr1, tr2, tr3], axis = 1).max(axis = 1)

        # 计算+DM和-DM
        plus_dm = high - high.shift()
        minus_dm = low.shift() - low
        mask_plus = (plus_dm > minus_dm) & (plus_dm > 0)
        mask_minus = (minus_dm > plus_dm) & (minus_dm > 0)
        plus_dm = plus_dm * mask_plus
        minus_dm = minus_dm * mask_minus

        # 计算+DI和-DI
        smooth_tr = tr.ewm(span=period).mean()
        smooth_plus_dm = plus_dm.ewm(span=period).mean()
        smooth_minus_dm = minus_dm.ewm(span=period).mean()
        plus_di = (smooth_plus_dm / smooth_tr) * 100
        minus_di = (smooth_minus_dm / smooth_tr) * 100

        # 计算ADX
        diff_di = abs(plus_di - minus_di)
        sum_di = plus_di + minus_di
        dx = (diff_di / sum_di) * 100
        adx = dx.ewm(span=period).mean()

        # 计算ADXR
        adxr = ((adx.shift(period) + adx) / 2)

        # +DI, -DI, ADX, ADXR
        return plus_di.values, minus_dm.values, adx.values, adxr.values

    def TRIX(self, period=30, m = 9):
        close = self.data['close']

        # 1. 计算N天的收盘价的指数平均AX
        ax = close.ewm(span=period, adjust=False).mean()

        # 2. 计算N天的AX的指数平均BX
        bx = ax.ewm(span=period, adjust=False).mean()

        # 3. 计算N天的BX的指数平均TRIX
        trix = bx.ewm(span=period, adjust=False).mean()

        # 计算TRIX的百分比变化率
        trix = trix.pct_change() * 100

        # 4. 计算TRIX的m日移动平均TRMA
        trma = trix.rolling(window=m).mean()

        # TRIX, TRMA
        return trix.values, trma.values

    def DMA(self, short_window = 10, long_window = 50):
        close = self.data['close']

        # 计算短期移动平均线
        short_sma = close.rolling(window=short_window).mean()

        # 计算长期移动平均线
        long_sma = close.rolling(window=long_window).mean()

        # 计算DMA
        dma = short_sma - long_sma

        # 返回DMA的Numpy数组形式，同时去除可能存在的NaN值
        return dma.values

    def BOLL(self, window=20, num_std=2):
        close = self.data['close']

        # 计算中间线（即移动平均线）
        middle_band = close.rolling(window=window).mean()

        # 计算标准差
        std_dev = close.rolling(window=window).std()

        # 计算上轨和下轨
        upper_band = middle_band + (std_dev * num_std)
        lower_band = middle_band - (std_dev * num_std)

        # 中轨线，上轨线，下轨线
        return middle_band.values, upper_band.values, lower_band.values

    def W_R(self, period=14):
        high_prices = self.data['high']
        low_prices = self.data['low']
        close_prices = self.data['close']

        highest_high = high_prices.rolling(window=period).max()
        lowest_low = low_prices.rolling(window=period).min()

        wr = ((highest_high - close_prices) / (highest_high - lowest_low)) * 100
        return wr.values

    def BIAS(self, period=10):
        # 计算N日简单移动平均
        ma_N = self.data['close'].rolling(window=period).mean()

        # 计算BIAS
        bias = ((self.data['close'] - ma_N) / ma_N) * 100

        # 将结果转换为NumPy数组并返回
        return bias.values

    def CR(self, period=10):
        high = self.data['high']
        low = self.data['low']

        #近period日的最高价与最低价的平均值
        high_avg = high.rolling(window=period).mean()
        low_avg = low.rolling(window=period).mean()

        #前一日的最高价与最低价
        recent_high = high.shift(1)
        recent_low = low.shift(1)

        #CR指标计算
        cr = (recent_high + recent_low) / 2
        cr = cr / low_avg * 100
        cr = cr / high_avg * low_avg
        return cr.values

    def VR(self,short_period=26,long_period=52):
        # 计算每日涨跌
        daily_change = self.data['close'] - self.data['pre_close']

        # 初始化累积成交量数组
        up_vol = np.zeros(len(self.data))
        down_vol = np.zeros(len(self.data))

        # 遍历数据，计算上涨和下跌日的成交量
        for i in range(1, len(self.data)):
            if daily_change.iloc[i] > 0 or (daily_change.iloc[i] == 0 and daily_change.iloc[i - 1] > 0):
                up_vol[i] = self.data['vol'].iloc[i]
            elif daily_change.iloc[i] < 0 or (daily_change.iloc[i] == 0 and daily_change.iloc[i - 1] < 0):
                down_vol[i] = self.data['vol'].iloc[i]

        # 计算短期和长期的平均成交量
        up_vol_short = pd.Series(up_vol).rolling(window=short_period).sum()
        down_vol_long = pd.Series(down_vol).rolling(window=long_period).sum()

        # 计算VR指标
        vr = up_vol_short / down_vol_long

        # 转换为NumPy数组并返回
        return vr.values

    def OBV(self):
        obv = np.zeros(len(self.data))

        # 设置初始OBV值为0
        obv[0] = 0

        # 遍历数据，计算OBV
        for i in range(1, len(self.data)):
            if self.data['close'].iloc[i] > self.data['close'].iloc[i - 1]:
                obv[i] = obv[i - 1] + self.data['vol'].iloc[i]
            elif self.data['close'].iloc[i] < self.data['close'].iloc[i - 1]:
                obv[i] = obv[i - 1] - self.data['vol'].iloc[i]
            else:
                obv[i] = obv[i - 1]

        return obv

    def MIKE(self, period=10):
        high = self.data['high']
        low = self.data['low']
        close = self.data['close']

        #计算起始价
        TYP = (high + low + close) / 3

        # period内的最高价和最低价
        maxx = high.rolling(window=period).max()
        minn = low.rolling(window=period).min()

        #计算初级压力WR，中级压力MR，强力压力SR
        WR = TYP * 2 - minn
        MR = TYP + (maxx - minn)
        SR = 2 * maxx - minn

        #计算初级支撑WS，中级支撑MS，强力支撑SS
        WS = TYP - (maxx - TYP)
        MS = TYP - (maxx - minn)
        SS = 2 * minn - maxx

        return np.stack([TYP, WR, MR, SR, WS, MS, SS], axis=1)

    def PVI(self, period=14):
        # 计算收盘价和成交量的移动平均
        avg_close = self.data['close'].rolling(window=period).mean()
        avg_vol = self.data['vol'].rolling(window=period).mean()

        # 初始化PVI值
        pvi = np.zeros(len(self.data))
        pvi[0] = 100.0  # 第一次计算时，PV以100代替

        # 计算PVI
        for i in range(1, len(self.data)):
            # 当前收盘价与过去N日收盘价的平均值的差异
            close_diff = self.data['close'].iloc[i] - avg_close.iloc[i]

            # 当前成交量与过去N日成交量的平均值的差异
            vol_diff = self.data['vol'].iloc[i] - avg_vol.iloc[i]

            # 如果成交量差大于0，则使用公式计算新的PVI值
            if vol_diff > 0:
                pvi[i] = pvi[i - 1] + (close_diff / avg_close.iloc[i]) * pvi[i - 1]
            else:
                pvi[i] = pvi[i - 1]

        return pvi

    def ASI(self):
        # 读取CSV文件
        df = self.data

        # 初始化ASI数组，长度与数据框行数相同，初始值为NaN
        asi = np.full(len(df), np.nan)

        # 如果数据框行数小于2，则直接返回全为NaN的数组
        if len(df) < 2:
            return asi

        # 计算ASI
        for i in range(1, len(df)):
            # 获取数据
            today_high = df['high'].iloc[i]
            today_low = df['low'].iloc[i]
            yesterday_close = df['close'].iloc[i - 1]
            yesterday_open = df['open'].iloc[i - 1]
            today_close = df['close'].iloc[i]
            today_open = df['open'].iloc[i]

            # 计算A、B、C、D、E、F、G
            A = abs(today_high - yesterday_close)
            B = abs(today_low - yesterday_close)
            C = abs(today_high - today_low)
            D = abs(yesterday_close - yesterday_open)
            E = today_close - yesterday_close
            F = today_close - today_open
            G = yesterday_close - yesterday_open

            # 计算X
            X = E + 0.5 * F + G

            # 比较A、B，取最大值
            maxx = max(A, B, C)
            if maxx == A:
                R = A + 0.5 * B + 0.25 * D
            elif maxx == B:
                R = B + 0.5 * A + 0.25 * D
            else:
                R = C + 0.25 * D

            # 计算K和L
            K = max(A, B)
            L = 13


            if(R == 0) :
                R = 1e-3
            if(L == 0) :
                L = 1e-3

            # 计算SI
            if R != 0 and L != 0:
                SI = 50 * (X / R) * (K / L)
            else:
                SI = np.nan

            # 计算ASI
            if i == 1:
                asi[i] = SI
            else:
                asi[i] = asi[i - 1] + SI

        return asi

    def KDJ(self, n=9):
        # 读取CSV文件到pandas DataFrame
        df = self.data

        # 初始化K、D、J值的数组，使用NaN填充
        K = np.full_like(df['close'], np.nan)
        D = np.full_like(df['close'], np.nan)
        J = np.full_like(df['close'], np.nan)

        # 计算RSV
        for i in range(n - 1, len(df)):
            # 计算过去n天内的最高价和最低价
            high = df['high'].iloc[i - n + 1:i + 1].max()
            low = df['low'].iloc[i - n + 1:i + 1].min()

            # 计算当日的RSV
            if high == low:
                RSV = 50  # 避免除零错误，当最高价等于最低价时，RSV设为50
            else:
                RSV = (df['close'].iloc[i] - low) / (high - low) * 100

            # 计算K值和D值
            if i == n - 1:
                # 如果是第一天，没有前一日的K值和D值，用50代替
                K[i] = RSV
                D[i] = K[i]
            else:
                K[i] = 2 / 3 * K[i - 1] + 1 / 3 * RSV
                D[i] = 2 / 3 * D[i - 1] + 1 / 3 * K[i]

            # 计算J值
            J[i] = 3 * D[i] - 2 * K[i]

        return np.stack([K, D, J], axis=1)

    def ROC(self, n=1):
        # 读取CSV文件到pandas DataFrame
        df = self.data

        # 初始化ROC数组，使用NaN填充
        ROC = np.full_like(df['close'], np.nan)

        # 计算ROC
        for i in range(n, len(df)):
            # 计算n天前的收盘价
            close_n_days_ago = df['close'].iloc[i - n]

            # 计算当日的ROC
            if close_n_days_ago == 0:
                close_n_days_ago = 1e-3

            ROC[i] = ((df['close'].iloc[i] - close_n_days_ago) / close_n_days_ago) * 100

        return ROC

    def MTM(self, n=3):
        # 读取CSV文件到pandas DataFrame
        df = self.data

        # 初始化ROC数组，使用NaN填充
        MTM = np.full_like(df['close'], np.nan)

        # 计算ROC
        for i in range(n, len(df)):
            MTM[i] = df['close'].iloc[i] - df['close'].iloc[i - n]

        return MTM

    def ARBR(self, n=26):
        # 读取CSV文件到pandas DataFrame
        df = self.data

        # 初始化AR和BR数组，使用NaN填充
        AR = np.full_like(df['close'], np.nan)
        BR = np.full_like(df['close'], np.nan)

        # 计算AR和BR
        for i in range(n - 1, len(df)):
            # 计算AR
            sum_HO = df['high'].iloc[i - n + 1:i + 1].sub(df['open'].iloc[i - n + 1:i + 1]).sum()
            sum_OL = df['open'].iloc[i - n + 1:i + 1].sub(df['low'].iloc[i - n + 1:i + 1]).sum()
            if sum_OL == 0:
                sum_OL = 1e-3
            #else:
            AR[i] = (sum_HO / sum_OL) * 100

            # 计算BR
            if i == 0:
                BR[i] = np.nan  # 第一天没有昨天的收盘价，无法计算BR
            else:
                sum_HPC = df['high'].iloc[i - n + 1:i + 1].sub(df['close'].iloc[i - 1]).sum()
                sum_PCL = df['close'].iloc[i - 1:i].sub(df['low'].iloc[i - n + 1:i + 1]).sum()
                if sum_PCL == 0:
                    sum_PCL = 1e-3
                #else:
                BR[i] = (sum_HPC / sum_PCL) * 100

        return np.stack([AR, BR], axis=1)

    def PSY(self, n=12):
        # 读取CSV文件到pandas DataFrame
        df = self.data

        # 初始化PSY数组，使用NaN填充
        PSY = np.full_like(df['close'], np.nan)

        # 计算PSY
        for i in range(n, len(df)):
            # 计算N日内上涨天数
            up_days_count = 0
            for j in range(i - n, i):
                if df['close'].iloc[j + 1] > df['close'].iloc[j]:
                    up_days_count += 1

            psy_value = up_days_count / n * 100
            PSY[i] = psy_value

        return PSY

    def SAR(self, n=4, initial_af=0.02, max_af=0.2):
        # 读取CSV数据到pandas DataFrame
        df = self.data

        # 检查数据是否足够计算SAR值
        if len(df) <= n:
            return np.full(len(df), np.nan)

        # 提取高点和低点
        highs = df['high'].to_numpy()
        lows = df['low'].to_numpy()

        # 初始化存储SAR值和加速因子的数组
        sar_values = np.full_like(highs, np.nan, dtype=float)
        acceleration_factors = np.full_like(highs, initial_af, dtype=float)

        # 确定初始SAR值 (SAR(T0))
        if highs[n] > highs[0]:
            sar_values[n] = lows[0]
        else:
            sar_values[n] = highs[0]

        # 确定EP (极点价) 的初始值
        if highs[n] > highs[0]:
            ep = max(highs[1:n + 1])
        else:
            ep = min(lows[1:n + 1])

        # 计算每个周期的SAR值
        for i in range(n + 1, len(df)):
            if highs[i - n:i].max() > ep:
                ep = highs[i - n:i].max()
            elif lows[i - n:i].min() < ep:
                ep = lows[i - n:i].min()

            if highs[i - 1] > highs[i - n:i - 1].max() and sar_values[i - 1] <= highs[i - 1]:
                sar_values[i] = highs[i - 1]
                acceleration_factors[i] = initial_af
            elif lows[i - 1] < lows[i - n:i - 1].min() and sar_values[i - 1] >= lows[i - 1]:
                sar_values[i] = lows[i - 1]
                acceleration_factors[i] = initial_af
            else:
                sar_values[i] = sar_values[i - 1] + acceleration_factors[i] * (ep - sar_values[i - 1])
                acceleration_factors[i] = min(acceleration_factors[i - 1] + initial_af, max_af)

        return sar_values

    def OBOS(self, n=5):
        # 读取CSV数据到pandas DataFrame
        df = self.data

        # 检查数据是否足够计算OBOS值
        if len(df) <= n:
            return np.full(len(df), np.nan)

        # 提取涨跌情况
        changes = df['close'].diff().dropna().to_numpy()

        # 初始化存储OBOS值的数组
        obos_values = np.full(len(df), np.nan)

        # 计算每日OBOS值
        for i in range(n, len(df)):
            # 统计N日内上涨和下跌家数
            na_sum = np.sum(changes[i - n + 1:i + 1] > 0)
            nb_sum = np.sum(changes[i - n + 1:i + 1] < 0)

            # 计算OBOS值
            obos_values[i] = na_sum - nb_sum

        return obos_values

    def CCI(self, n=20):
        # 读取CSV数据到pandas DataFrame
        df = self.data

        # 检查数据是否足够计算CCI值
        if len(df) <= n:
            return np.full(len(df), np.nan)

        # 计算典型价格（TP）
        tp = (df['high'] + df['low'] + df['close']) / 3

        # 计算MA（N日收盘价的简单移动平均）
        ma = tp.rolling(window=n, min_periods=1).mean()

        # 计算MD（N日的平均绝对偏差）
        md = np.abs(ma - df['close']).rolling(window=n, min_periods=1).mean()

        # 计算CCI值
        cci = (tp - ma) / (0.015 * md)

        return cci.to_numpy()

    def WVAD(self, n=24, m=6):
        # 读取CSV数据到pandas DataFrame
        df = self.data

        # 检查数据是否足够计算WVAD和MAWVAD
        if len(df) <= n:
            return np.full(len(df), np.nan)

        # 计算A和B
        A = df['close'] - df['open']
        B = df['high'] - df['low']

        # 计算C
        C = A / B * df['vol']

        # 计算WVAD
        WVAD = C.rolling(window=n, min_periods=1).sum()

        # 计算MAWVAD
        MAWVAD = WVAD.rolling(window=m, min_periods=1).mean()

        return MAWVAD.to_numpy()

    def TAPI(self):
        # 读取CSV数据到pandas DataFrame
        df = self.data

        tapi_values = np.full(len(df), np.nan)

        # 计算每日OBOS值
        for i in range(1, len(df)):
            # 统计N日内上涨和下跌家数
            a = df['amount'].iloc[i]
            b = df['amount'].iloc[i - 1]

            if (b == 0):
                tapi_values[i] = 1e3

            tapi_values[i] = a + a / b * 100

        return tapi_values

"""
将多个股票序列作为数据集
输入的ts_codes是一个股票代码序列:如['000001.SZ', '000006.SZ', .... ,'000005.SZ']
"""
def StackedDataSet(ts_codes, start_date = None, end_date = None):
    stackedX_list = []
    stackedY_list = []
    for ts_code in ts_codes:
        dataset = DataSet(ts_code, start_date, end_date).data
        # 某一只股票的特征与标签
        X = dataset.drop(columns='Fmark')[:-1].values
        y = dataset['Fmark'][1:].values
        stackedX_list.append(X)
        stackedY_list.append(y)

    stackedX = np.concatenate(stackedX_list, axis=0)
    stackedY = np.concatenate(stackedY_list, axis=0).astype('int64')
    return stackedX, stackedY

if __name__ == "__main__":

    pd.set_option('display.max_rows', None)  # 显示所有行
    pd.set_option('display.max_columns', None)  # 显示所有列
    pd.set_option('display.width', None)  # 自动调整宽度，避免换行
    pd.set_option('display.max_colwidth', -1)  # 不限制列宽，显示完整内容
    np.set_printoptions(threshold=np.inf)
    pull = pullDataFromTushare()

    #stock = '000001.SZ'
    #dataset = DataSet(stock, '20200201', '20210201')
    ##dataset.data.to_csv(stock+'.csv', index=False)
    #print(dataset.data)
    #print(dataset.data['Fmark'])
    print(StackedDataSet(['000001.SZ', '000006.SZ'], '20200201', '20210201'))